# coding: utf8

"""

TAO Ain't Ollydbg

TAO(This Ain't Ollydbg) is an open-source OllyDBG parody in Python + Qt.

"""


from qt import *
from dockablewidgets import *


from winappdbg import HexDump, Debug, System, Process
from winappdbg import EventHandler

import winappdbg
import dragndrop


import os
import sys
import time
import logging
import traceback

from tablewindow import TableWindow
from cpuwindow import CpuWindow
from logwindow import LogWindow, log
from processwindow import ProcessWindow
from modulewindow import ModuleWindow
from threadwindow import ThreadWindow
from optionwindow import OptionWindow
from aboutwindow import AboutWindow






class DebugEventHandler(EventHandler):

    def exception(self, event):
        thread = event.get_thread()
        process = event.get_process()
        pc = thread.get_pc()
        log('Exception happened at %08X!' % pc)
        sp = thread.get_sp()
        ctx = thread.get_context()
        cpuwindow = self.window.cpuwindow

        cpuwindow.show_disasm(process, thread, pc)
        cpuwindow.show_dump(process, pc)
        cpuwindow.show_stack(process, sp)
        cpuwindow.show_reg(ctx)


        return 'pause'

    def default(self, event):
        log(event)
        name = event.get_event_name()
        code = event.get_event_code()
        pid = event.get_pid()
        tid = event.get_tid()
        pc = event.get_thread().get_pc() # pc = program counter = EIP

        format_string = "%s (%s) at address %s, process %d, thread %d"
        message = format_string % ( name, HexDump.integer(code), HexDump.address(pc), pid, tid )
        log(message)

    def load_dll(self, event):
        self.default(event)
        self.window.modulewindow.update_modules(event.get_process())
        module = event.get_module()
        log(module)

    create_process = create_thread =  default

    def __init__(self, window):
        super(DebugEventHandler, self).__init__()
        self.window = window

class DebugStatus(object):
    IDLE, WAIT, PAUSED, STEP_INTO, STEP_OVER, STOP = range(6)

class DebugThread(QThread):

    def __init__(self, window):
        super(DebugThread, self).__init__()
        self.window = window

    def run(self):
        log('Run %s' % self.window.filename)
        self.debug_loop()

    def debug_loop(self):
        eventhandler = DebugEventHandler(window=self.window)
        debug = self.debug = Debug(eventHandler=eventhandler, bKillOnExit=True)
        debug.execl(self.window.filename)
        self.status = 'wait'
        try:
            while debug.get_debugee_count() > 0:
                if self.status == 'wait':
                    event = debug.wait()
                    self.status = debug.dispatch(event) or 'continue'

                elif self.status == 'pause':
                    1


                elif self.status == 'step_into':
                    debug.start_tracing(event.get_tid())
                    self.cont()
                    
                elif self.status == 'continue':
                    self.status = 'wait'
                    debug.cont(event)
                    
                elif self.status == 'stop':
                    break

                time.sleep(0)

        finally:
            debug.stop()

        self.status = 'idle'

    def pause(self):
        self.status = 'pause'

    def step_into(self):
        self.status = 'step_into'
        
    def step_over(self):
        1/0

    def cont(self):
        self.status = 'continue'


    

class MainWindow(QMainWindow):
    TITLE = "TAO Ain't Ollydbg"

    def __init__(self, parent=None):
        super(MainWindow, self).__init__(parent)

        self.setWindowTitle(self.TITLE)

        self.create_tab()
        self.create_menubar()
        self.create_toolbar()
        self.create_cmdbar()
        self.create_statusbar()
        
        self.tab.setCornerWidget(self.toolbar)

        dragndrop.setup_drag_and_drop(self, self.open_app)
        
        self.showMaximized()

    def create_tab(self):
        tab = self.tab = QTabWidget()
        
        self.tabs = []
        for cls in [CpuWindow, LogWindow, ProcessWindow, ModuleWindow, ThreadWindow, OptionWindow, AboutWindow]:
            window = cls(self)
            setattr(self, cls.__name__.lower(), window)
            tab.addTab(window, window.windowTitle())
            self.tabs.append(window)

        tab.setMovable(True)
        tab.setTabsClosable(True)
        tab.tabCloseRequested.connect(self.close_tab)

        self.setCentralWidget(tab)

    def close_tab(self, index):
        self.tab.removeTab(index)


    def restart_debuggee(self):
        self.debugthread.status = 'stop'
        while self.debugthread.status != 'idle':
            time.sleep(0)
        self.debugthread.terminate()
        del self.debugthread # throw it to GC
        self.open_app(self.filename)
        
    def create_menubar(self):
        self.menubar = menubar = QDockableMenuBar(self)

        file = menubar.addMenu('&File')
        file.addAction(QAction('Open', self, shortcut='F3', statusTip='Open debuggee', triggered=self.on_open))
        file.addAction(QAction('Test:notepad', self, shortcut='Ctrl+1', triggered=lambda: self.open_app('C:/windows/system32/notepad.exe')))
        file.addSeparator()
        file.addAction(QAction('Restart', self, shortcut='Ctrl+F2', triggered=self.restart_debuggee))
        file.addAction(QAction('Exit', self, shortcut='Alt+X', statusTip='Close debugger', triggered=self.close))

        view = menubar.addMenu('&View')

        def make_show(tab):
            def f():
                self.show_tab(tab)
            return f

        for tab in self.tabs:
            title = str(tab.windowTitle())
            action = QAction('&%s'%title, self, shortcut='Alt+%s'%title[0].upper(), triggered=make_show(tab))
            view.addAction(action)
            
        view.addSeparator()
        view.addAction(QAction('Disassemble backward', self, shortcut='Ctrl+Up'))
        view.addAction(QAction('Disassemble forward', self, shortcut='Ctrl+Down'))

        debug = menubar.addMenu('&Debug')
        debug.addAction(QAction('&Run', self, shortcut='F9', triggered=lambda : self.debugthread.cont()))
        debug.addAction(QAction('Run &Thread', self, shortcut='F11'))
        debug.addSeparator()
        debug.addAction(QAction('Step &Into', self, shortcut='F7', triggered=lambda : self.debugthread.step_into()))
        debug.addAction(QAction('Step &Over', self, shortcut='F8', triggered=lambda : self.debugthread.step_over()))
        debug.addAction(QAction('Execute &Until Return', self, shortcut='Ctrl+F9'))

        trace = menubar.addMenu('&Trace')
        
        plugins = menubar.addMenu('&Plugins')
        
        options = menubar.addMenu('&Options')
        options.addAction(QAction('&Options', self))
        
        windows = menubar.addMenu('&Windows')        
        windows.addAction(QAction('&Always On Top', self))
        
        help = menubar.addMenu('?')
        help.addAction(QAction('About', self, triggered=lambda: self.show_tab(self.about)))

    def show_tab(self, tabwin):
        index = self.tab.indexOf(tabwin)

        if index == -1:
            current = self.tab.currentIndex()
            self.tab.insertTab(current+1, tabwin, tabwin.windowTitle())

        self.tab.setCurrentWidget(tabwin)

    def create_toolbar(self):
        toolbar = self.toolbar = QToolBar(self)
        self.addToolBar(toolbar)
        self.toolbar.setIconSize(QSize(16, 16))
        
        for name in 'nan wu e mi tuo fo'.split():
            action = QAction(QIcon('icons/%s.bmp'%name), name, self)
            toolbar.addAction(action)        

    def create_cmdbar(self):
        cmdbar = QToolBar(self)
        self.addToolBar(Qt.BottomToolBarArea, cmdbar)

        lbl = QLabel(' Command ')
        line = QLineEdit(cmdbar)
        cmdbar.addWidget(lbl)
        cmdbar.addWidget(line)

    def create_statusbar(self):
        self.statusbar = QDockableStatusBar(self)
        self.statusbar.showMessage('Ready')

    def open_app(self, path):
        self.setWindowTitle(path)
        self.filename = str(path)
        debugthread = self.debugthread = DebugThread(self)
        debugthread.start()

    def on_open(self):
        filename = QFileDialog.getOpenFileName(self, 'Open file')
        if filename:
            self.open_app(filename)

def main():
    System.request_debug_privileges()

    app = QApplication(sys.argv)
    wnd = MainWindow()
    wnd.show()
    sys.exit(app.exec_())


if __name__ == '__main__':
    main()

